import os
import shutil
from sklearn.model_selection import train_test_split


def split_dataset(images_dir, labels_dir, output_image_dir, output_label_dir, test_size=0.2):
    """
    拆分YOLO数据集为训练集和验证集。

    参数：
    images_dir (str): 包含所有图像文件的目录路径。
    labels_dir (str): 包含所有标签文件的目录路径。
    output_train_dir (str): 输出训练集图像和标签的目录路径。
    output_val_dir (str): 输出验证集图像和标签的目录路径。
    test_size (float): 验证集的比例，范围在0到1之间。
    """
    # 确保输出目录存在
    os.makedirs(output_image_dir, exist_ok=True)
    os.makedirs(output_label_dir, exist_ok=True)

    # 获取所有图像和标签文件名
    image_files = [f for f in os.listdir(images_dir) if f.endswith('.jpg') or f.endswith('.png')]
    label_files = [f for f in os.listdir(labels_dir) if f.endswith('.txt')]

    # 确保图像和标签数量匹配
    assert len(image_files) == len(label_files), "图像和标签数量不匹配"

    # 拆分数据集
    train_images, val_images, train_labels, val_labels = train_test_split(image_files, label_files, test_size=test_size)

    # 复制训练集图像和标签到输出目录
    for train, val in zip(train_images, val_images):
        shutil.copy(os.path.join(images_dir, train), os.path.join(output_image_dir+'/train', train))
        shutil.copy(os.path.join(images_dir, val), os.path.join(output_image_dir+'/val', val))

    # 复制验证集图像和标签到输出目录
    for train, val in zip(train_labels, val_labels):
        shutil.copy(os.path.join(images_dir, train), os.path.join(output_label_dir+'/train', train))
        shutil.copy(os.path.join(labels_dir, val), os.path.join(output_label_dir+'/val', val))


if __name__ == "__main__":
    images_dir = r"C:\Users\26954\Desktop\label"
    labels_dir = r"C:\Users\26954\Desktop\label"
    output_image_dir = r"C:\Users\26954\Desktop\label_result\images"
    output_label_dir = r"C:\Users\26954\Desktop\label_result\labels"
    split_dataset(images_dir, labels_dir, output_image_dir, output_label_dir)
